package com.nx.platform.es.bean.modle.rescore;

import com.google.common.collect.ImmutableMap;
import com.nx.platform.es.bean.modle.score.ScoreFunctionField;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.collections4.MapUtils;
import org.elasticsearch.common.lucene.search.function.CombineFunction;
import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder.FilterFunctionBuilder;
import org.elasticsearch.search.rescore.QueryRescorerBuilder;
import org.elasticsearch.search.rescore.RescorerBuilder;

import java.util.LinkedList;
import java.util.List;
import java.util.Map;

/**
 * @author
 * @since 2016年10月16日
 */
public class FunctionScoresHandler implements RescoreFieldHandler {

    @Override
    public RescorerBuilder<?> handle(ImmutableMap<String, ?> fieldConfig, List<ScoreFunctionField> functions,
            Map<String, Object> params) {
        if (CollectionUtils.isEmpty(functions)) {
            return null;
        }

        List<FilterFunctionBuilder> fs = new LinkedList<>();
        for (ScoreFunctionField scoreFunctionField : functions) {
            FilterFunctionBuilder f = scoreFunctionField.handle(params);
            if (f != null) {
                fs.add(f);
            }
        }
        if (fs.isEmpty()) {
            return null;
        }
        FunctionScoreQueryBuilder queryBulder = QueryBuilders.functionScoreQuery(
                fs.toArray(new FilterFunctionBuilder[fs.size()]));
        String boostMode = MapUtils.getString(fieldConfig, "boost_mode", "sum");
        String scoreMode = MapUtils.getString(fieldConfig, "score_mode", "sum");
        queryBulder.boostMode(CombineFunction.fromString(boostMode))
                .scoreMode(FunctionScoreQuery.ScoreMode.fromString(scoreMode));

        QueryRescorerBuilder rescorer = new QueryRescorerBuilder(queryBulder);
        float queryWeight = MapUtils.getFloatValue(fieldConfig, "query_weight", 1f);
        if (Float.floatToIntBits(queryWeight) != 1) {
            rescorer.setQueryWeight(queryWeight);
        }
        float rescoreQueryWeight = MapUtils.getFloatValue(fieldConfig, "rescore_query_weight", 1f);
        if (Float.floatToIntBits(rescoreQueryWeight) != 1) {
            rescorer.setRescoreQueryWeight(rescoreQueryWeight);
        }
        return rescorer;
    }

}
